import {BaseChatModel} from 'langchain/chat_models/base'
import {AIMessage} from 'langchain/schema'

export class CheeseChatModel extends BaseChatModel {
    constructor(fields) {
        super(fields ?? {});
        Object.defineProperty(this, "streaming", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        Object.defineProperty(this, "baseUrl", {
            enumerable: true,
            configurable: true,
            writable: true,
            value: false
        });
        this.baseUrl = fields?.baseUrl ?? 'http://43.134.44.205:30625/context_batch_return';
        this.streaming = fields?.streaming ?? false;
    }
    _llmType() {
        return "cheese_chat";
    }
    async _call(messages, options, runManager) {
        const history = []
        const is_bots = []
        messages.forEach(mess=>{
            history.push(mess.content)
            switch (mess._getType()){
                case "ai":
                    is_bots.push(1)
                    break
                case "human":
                    is_bots.push(0)
                    break
            }
        })
        is_bots.pop()
        const prompt = history.pop()
        const baseUrl = this.baseUrl
        const headers = {
            "Content-Type": "application/json",
        };
        const data = {
            text: prompt,
            message: history,
            is_bots: is_bots
        };
        return await this.caller.callWithOptions({signal: options.signal}, async () => {
            const response = await fetch(baseUrl, {
                method: "POST",
                headers,
                body: JSON.stringify(data),
                signal: options.signal,
            });
            const reader = response.body.getReader();
            const decoder = new TextDecoder('utf-8');
            let last_index = 0;
            let buffer = '';
            let res = '';
            while (true) {
                try {
                    const {done, value} = await reader.read()
                    if (done) {
                        break;
                    }
                    buffer += decoder.decode(value);
                    const delimiterIndex = buffer.indexOf('\0');
                    if (delimiterIndex !== -1) {
                        let decodedText = buffer.substring(0, delimiterIndex);
                        buffer = '';
                        try {
                            decodedText = JSON.parse(decodedText);
                            if (decodedText.status === 'Success') {
                                res = decodedText.text
                                if(this.streaming){
                                    void runManager?.handleLLMNewToken(decodedText.text.slice(last_index))
                                }
                                last_index = decodedText.text.length;
                            } else if (decodedText.status === 'Fail') {
                                break;
                            }
                        } catch (e) {
                            break
                        }
                    }
                } catch (e) {
                    break
                }
            }
            return res
        })
    }
    async _generate(messages, options, runManager) {
        const text = await this._call(messages, options, runManager);
        const message = new AIMessage(text);
        return {
            generations: [
                {
                    text: message.content,
                    message,
                },
            ],
        };
    }
}
